package com.yololabelmaker;

import javax.swing.*;
import java.awt.image.BufferedImage;
import java.io.*;
import java.nio.file.*;
import java.util.List;

public class YOLOLabelSaver {
    // 类别映射表，用于存储类别名称和ID的对应关系
    private static java.util.Map<String, Integer> classMap = new java.util.HashMap<>();
    private static int nextClassId = 0;
    
    static {
        // 尝试加载已有的类别映射文件
        loadClassMap();
    }
    
    // 保存标签
    public void saveLabels(String imagePath, BufferedImage image, List<LabelRect> labelRects) {
        if (imagePath == null || image == null || labelRects.isEmpty()) {
            return;
        }
        
        // 确保标签目录存在
        File imageFile = new File(imagePath);
        String labelsDirPath;
        
        // 检查是否是网络图片
        if (imagePath.startsWith("network_image_")) {
            // 对于网络图片，保存到当前目录的labels文件夹
            labelsDirPath = System.getProperty("user.dir") + File.separator + "labels";
        } else {
            // 对于本地图片，按照YOLO数据集结构保存
            String parentDir = imageFile.getParent();
            if (parentDir != null && parentDir.endsWith("images")) {
                // 如果图片在images文件夹中，则标签应该放在同级的labels文件夹中
                labelsDirPath = parentDir.replace("images", "labels");
            } else {
                // 否则，在图片所在目录创建labels文件夹
                labelsDirPath = parentDir + File.separator + "labels";
            }
        }
        
        // 创建labels目录
        File labelsDir = new File(labelsDirPath);
        if (!labelsDir.exists()) {
            labelsDir.mkdirs();
        }
        
        // 生成标签文件名（与图片文件名相同，只是扩展名改为txt）
        String imageFileName = imageFile.getName();
        String labelFileName = imageFileName.substring(0, imageFileName.lastIndexOf('.')) + ".txt";
        String labelFilePath = labelsDirPath + File.separator + labelFileName;
        
        try (BufferedWriter writer = new BufferedWriter(new FileWriter(labelFilePath))) {
            int imageWidth = image.getWidth();
            int imageHeight = image.getHeight();
            
            // 写入每个标签框
            for (LabelRect rect : labelRects) {
                // 获取或分配类别ID
                int classId = getClassId(rect.getClassName());
                
                // 计算YOLO格式的坐标和尺寸
                double centerX = rect.getYOLOCenterX(imageWidth);
                double centerY = rect.getYOLOCenterY(imageHeight);
                double width = rect.getYOLOWidth(imageWidth);
                double height = rect.getYOLOHeight(imageHeight);
                
                // 写入YOLO格式的行
                writer.write(String.format("%d %.6f %.6f %.6f %.6f\n", 
                                          classId, centerX, centerY, width, height));
            }
            
            // 保存类别映射文件
            saveClassMap();
            
            // 显示保存成功信息
            JOptionPane.showMessageDialog(null, "标签保存成功：\n" + labelFilePath);
        } catch (IOException e) {
            JOptionPane.showMessageDialog(null, "保存标签时发生错误: " + e.getMessage());
            e.printStackTrace();
        }
    }
    
    // 获取类别ID，如果不存在则分配一个新的
    private int getClassId(String className) {
        if (!classMap.containsKey(className)) {
            classMap.put(className, nextClassId++);
        }
        return classMap.get(className);
    }
    
    // 保存类别映射文件
    private static void saveClassMap() {
        String classMapFilePath = System.getProperty("user.dir") + File.separator + "classes.txt";
        try (BufferedWriter writer = new BufferedWriter(new FileWriter(classMapFilePath))) {
            // 创建一个按ID排序的列表
            java.util.List<java.util.Map.Entry<String, Integer>> sortedEntries = 
                new java.util.ArrayList<>(classMap.entrySet());
            sortedEntries.sort(java.util.Map.Entry.comparingByValue());
            
            // 写入每个类别
            for (java.util.Map.Entry<String, Integer> entry : sortedEntries) {
                writer.write(entry.getKey() + "\n");
            }
        } catch (IOException e) {
            System.err.println("保存类别映射文件时发生错误: " + e.getMessage());
        }
    }
    
    // 加载类别映射文件
    private static void loadClassMap() {
        String classMapFilePath = System.getProperty("user.dir") + File.separator + "classes.txt";
        File classMapFile = new File(classMapFilePath);
        if (classMapFile.exists()) {
            try (BufferedReader reader = new BufferedReader(new FileReader(classMapFile))) {
                String line;
                int id = 0;
                while ((line = reader.readLine()) != null) {
                    line = line.trim();
                    if (!line.isEmpty()) {
                        classMap.put(line, id++);
                        nextClassId = id;
                    }
                }
            } catch (IOException e) {
                System.err.println("加载类别映射文件时发生错误: " + e.getMessage());
            }
        }
    }
}